
#!/usr/bin/env python3
import time
import struct
import argparse
import serial
import pilink
import pilink_parser

PILINK_TYPE_TERMINAL = 0x0004  # from pilink.py



WORD_SIZE = 4
MAX_RETRY = 5

class BootloaderFlasher:
    def __init__(self, port, baudrate):
        self.parser = pilink_parser.pilink_serial_parser()
        self.parser.set_msg_callback(self.on_msg)
        self.parser.set_port(port, 0.1)
        self.bootloader_ready = False
        self.new_ack = None
        self.new_cmd = None
        self.new_data = None
        self.new_ter = None

    def on_msg(self, port, buf, length):
        # Handle ACK messages
        if port == pilink_parser.pilink.PILINK_TYPE_BOOT_ACK:
            rack = pilink_parser.pilink.pilink_boot_ack_decode(buf)
            self.new_ack = rack
        # Handle terminal prompt messages
        elif port ==  pilink_parser.pilink.PILINK_TYPE_BOOT_CMD:
            rcmd = pilink_parser.pilink.pilink_boot_cmd_decode(buf)
            self.new_cmd = rcmd
        elif port == pilink_parser.pilink.PILINK_TYPE_BOOT_DATA:
            rdat = pilink_parser.pilink.pilink_boot_data_decode(buf)
            self.new_data = rdat
        elif port == pilink_parser.pilink.PILINK_TYPE_TERMINAL:
            rter = pilink_parser.pilink.pilink_terminal_data_decode(buf)
            self.new_ter = rter


    def trigger_bootloader(self, timeout=10):
        print("[STEP] Waiting for bootloader prompt...")
        deadline = time.time() + timeout
        while time.time() < deadline:
            self.parser.send_boot_cmd(pilink_parser.pilink.PILINK_BOOT_CMD_ABORT, 0, 0, 0)
            self.new_ter = None
            self.parser.update()
            if self.new_ter != None:
                if b"Aborted" in self.new_ter.msg:
                    return True
            time.sleep(0.1)
        return False
    
    def erase_flash(self, timeout=10):
        print("[STEP] Erase...")
        deadline = time.time() + timeout
        while time.time() < deadline:
            self.parser.send_boot_cmd(pilink_parser.pilink.PILINK_BOOT_CMD_ERASE, 0, 0, 0)
            self.new_ter = None
            self.parser.update()
            if self.new_ter != None:
                if b"Erased" in self.new_ter.msg:
                    return True
            time.sleep(0.1)
        return False

    def trigger_upgrade(self, timeout=10):
        print("[STEP] Start Upgrading...")
        deadline = time.time() + timeout
        while time.time() < deadline:
            self.parser.send_boot_cmd(pilink_parser.pilink.PILINK_BOOT_CMD_STA_UGD, 0, 0, 0)
            self.new_ter = None
            self.parser.update()
            if self.new_ter != None:
                if b"Upgrading" in self.new_ter.msg:
                    return True
            time.sleep(0.1)
        return False
    
    def end_upgrade(self, timeout=10):
        print("[STEP] Stop Upgrading...")
        deadline = time.time() + timeout
        while time.time() < deadline:
            self.parser.send_boot_cmd(pilink_parser.pilink.PILINK_BOOT_CMD_END_UGD, 0, 0, 0)
            self.new_ter = None
            self.parser.update()
            if self.new_ter != None:
                if b"Upgraded" in self.new_ter.msg:
                    return True
            time.sleep(0.1)
        return False
    
    def restart(self, timeout=10):
        print("[STEP] Restart")
        deadline = time.time() + timeout
        while time.time() < deadline:
            self.parser.send_boot_cmd(pilink_parser.pilink.PILINK_BOOT_CMD_RESTART, 0, 0, 0)
            self.new_ter = None
            self.parser.update()
            if self.new_ter != None:
                if b"Restarting" in self.new_ter.msg:
                    return True
            time.sleep(0.1)
        return False


    def wait_for_ack(self, addr, expected):
        timeout = 0.5
        deadline = time.time() + timeout
        while time.time() < deadline:
            time.sleep(0.001)
            self.new_ack = None
            self.parser.update()
            if self.new_ack != None:
                if self.new_ack.addr == addr and self.new_ack.data == expected:
                    return True
        return False

    def flash_firmware(self, firmware_path, base_addr):
        with open(firmware_path, 'rb') as f:
            data = f.read()
        if len(data) % WORD_SIZE:
            data += b'\xFF' * (WORD_SIZE - len(data) % WORD_SIZE)

        # Step 1: wait for bootloader
        if not self.trigger_bootloader():
            print("[ERROR] Bootloader not detected. Please reset into bootloader mode.")
            return False
        print("Bootloader detected.")
        
        # Step 2: erase
        if not self.erase_flash():
            print("[ERROR] Failed to erase flash.")
            return False
        print("Erasing completed.")

        # Step 3: start upgrade session and wait ack
        if not self.trigger_upgrade():
            print("[ERROR] Failed to start upgradung.")
            return False
        print("Upgrading started.")

        # Step 4: send data words
        print("[STEP] Sending firmware data")
        data_len = len(data)
        for offset in range(0, data_len, WORD_SIZE):
            word = struct.unpack('<I', data[offset:offset+4])[0]
            addr = base_addr + offset
            pct = offset * 100 / data_len
            print(f"\rProgramming:{pct:.2f}%", end='', flush=True)
            for attempt in range(1, MAX_RETRY+1):
                self.parser.send_boot_data(addr, word)
                if self.wait_for_ack(addr, word):
                    break
                print(f"[WARN] Retry {attempt} for address 0x{addr:08X}")
            else:
                print(f"[ERROR] Failed to program at address 0x{addr:08X}")
                return False

        # Step 5: finalize and restart
        print("[STEP] Finalizing upgrade")
        if not self.end_upgrade():
            print("[ERROR] Failed to stop upgradung.")
            return False
        print("Upgrading finished.")

        print("[STEP] Restarting target system")
        if not self.restart():
            print("[ERROR] Failed to restart system")
            return False
        
        print("Upgrading successfully.")

        return True

def main():
    parser = argparse.ArgumentParser(description="PiLink Bootloader Firmware Flasher")
    parser.add_argument('--port', required=True, help='Serial port, e.g., COM5 or /dev/ttyUSB0')
    parser.add_argument('--baudrate', type=int, default=115200, help='Serial baudrate')
    parser.add_argument('--file', required=True, help='Path to firmware .bin file')
    parser.add_argument('--base', type=lambda x: int(x, 0), default=0x08008000,
                        help='Base flash start address (hex or dec)')
    args = parser.parse_args()

    flasher = BootloaderFlasher(args.port, args.baudrate)
    success = flasher.flash_firmware(args.file, args.base)
    if success:
        print("[DONE] Firmware upgrade completed successfully.")
    else:
        print("[FAIL] Firmware upgrade failed.")

if __name__ == '__main__':
    main()
